import asyncio
from typing import Optional
from contextlib import AsyncExitStack

from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client

from anthropic import Anthropic
from dotenv import load_dotenv

from langchain_anthropic import ChatAnthropic
from langchain.tools import Tool, StructuredTool
from langchain.agents import  AgentExecutor, create_openai_tools_agent, create_structured_chat_agent
from langchain.prompts import ChatPromptTemplate
from langchain_core.messages import HumanMessage, SystemMessage

from services.agent_service import  create_llm, create_agent

load_dotenv()  # load environment variables from .env

class MCPClient:
    def __init__(self):
        # Initialize session and client objects
        self.session: Optional[ClientSession] = None
        self.exit_stack = AsyncExitStack()
        # self.anthropic = Anthropic()
        self.llm = create_llm()

    async def connect_to_server(self, server_script_path: str):
        """Connect to an MCP server
        
        Args:
            server_script_path: Path to the server script (.py or .js)
        """
        is_python = server_script_path.endswith('.py')
        is_js = server_script_path.endswith('.js')
        if not (is_python or is_js):
            raise ValueError("Server script must be a .py or .js file")
            
        command = "python3" if is_python else "node"
        server_params = StdioServerParameters(
            command=command,
            args=[server_script_path],
            env=None
        )
        
        stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
        self.stdio, self.write = stdio_transport
        self.session = await self.exit_stack.enter_async_context(ClientSession(self.stdio, self.write))
        
        await self.session.initialize()
        
        # List available tools
        response = await self.session.list_tools()
        tools = response.tools
        print("\nConnected to server with tools:", [tool.name for tool in tools])

    async def process_query(self, query: str) -> str:
        """Process a query using Claude and available tools"""


        response = await self.session.list_tools()
        # 将 MCP 工具转换为 LangChain 工具
        tools = []
        for tool in response.tools:
            print(f"Tool {tool.name}: {tool.description}: {tool.inputSchema}")

            # 将 JSON Schema 转换为 Pydantic 模型
            try:
                from pydantic import create_model
                import json
                
                # 解析 JSON Schema
                if isinstance(tool.inputSchema, str):
                    schema_dict = json.loads(tool.inputSchema)
                else:
                    schema_dict = tool.inputSchema
                
                # 创建字段定义
                fields = {}
                properties = schema_dict.get("properties", {})
                required = schema_dict.get("required", [])
                
                for prop_name, prop_schema in properties.items():
                    # 确定字段类型
                    field_type = str  # 默认为字符串类型
                    if prop_schema.get("type") == "number":
                        field_type = float
                    elif prop_schema.get("type") == "integer":
                        field_type = int
                    elif prop_schema.get("type") == "boolean":
                        field_type = bool
                    
                    # 确定是否为必填字段
                    is_required = prop_name in required
                    default = ... if is_required else None
                    
                    fields[prop_name] = (field_type, default)
                
                # 创建 Pydantic 模型
                args_schema = create_model(f"{tool.name}Schema", **fields)
                
            except Exception as e:
                print(f"无法为工具 {tool.name} 创建 Pydantic 模型: {str(e)}")
                # 如果转换失败，使用默认的字典类型
                from pydantic import BaseModel, Field
                class DefaultSchema(BaseModel):
                    args: dict = Field(description=f"工具 {tool.name} 的参数")
                args_schema = DefaultSchema

            # async def tool_func(tool_input: dict, tool_name=tool.name):
            #     print(f"Calling tool {tool_name} with input {tool_input}")
            #     # 确保输入是字典类型
            #     if not isinstance(tool_input, dict):
            #         raise ValueError(f"Tool {tool_name} expects a dictionary input")
            #     result = await self.session.call_tool(tool_name, tool_input)
            #     return result.content
            # 修改工具函数定义，使用 **kwargs 接收任意关键字参数
            async def tool_func(**kwargs):
                tool_name = tool.name  # 使用闭包捕获当前循环的 tool.name
                print(f"Calling tool {tool_name} with input {kwargs}")
                try:
                    result = await self.session.call_tool(tool_name, kwargs)
                    return "工具执行成功"
                    if not result.content:
                        return "工具执行成功，但没有返回内容"
                    if isinstance(result.content, (dict, list)):
                        # 如果返回的是结构化数据，转换为字符串
                        return json.dumps(result.content, ensure_ascii=False).encode('utf-8')
                    return str(result.content)  # 确保返回字符串类型
                except Exception as e:
                    return f"工具执行出错: {str(e)}"

            tools.append(
                StructuredTool(
                    name=tool.name,
                    description=tool.description,
                    func=tool_func,
                    args_schema=args_schema
                )
            )
        # 创建代理
        # prompt = ChatPromptTemplate.from_messages([
        #     ("system", "你是一个有用的助手，可以使用提供的工具来完成任务。请仔细分析用户需求，合理使用工具。"),
        #     ("human", "{input}"),
        #     ("assistant", "{agent_scratchpad}") # 这是一个占位符，用于在代理执行时插入工具调用结果
        # ])

        agent = create_agent(tools, self.llm)
        # agent = create_structured_chat_agent(self.llm, tools, prompt)
        # agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

        # 现在有一个会计凭证编号为1234，帮我生成摘要并返回文件路径；用中文回答
        # 执行代理
        result = await agent.ainvoke({"input": query})
        return result["output"]


        available_tools = [{ 
            "name": tool.name,
            "description": tool.description,
            "input_schema": tool.inputSchema
        } for tool in response.tools]

        # Initial Claude API call
        response = self.anthropic.messages.create(
            model="claude-3-5-sonnet-20241022",
            max_tokens=1000,
            messages=messages,
            tools=available_tools
        )

        # Process response and handle tool calls
        tool_results = []
        final_text = []

        for content in response.content:
            if content.type == 'text':
                final_text.append(content.text)
            elif content.type == 'tool_use':
                tool_name = content.name
                tool_args = content.input
                
                # Execute tool call
                result = await self.session.call_tool(tool_name, tool_args)
                tool_results.append({"call": tool_name, "result": result})
                final_text.append(f"[Calling tool {tool_name} with args {tool_args}]")

                # Continue conversation with tool results
                if hasattr(content, 'text') and content.text:
                    messages.append({
                      "role": "assistant",
                      "content": content.text
                    })
                messages.append({
                    "role": "user", 
                    "content": result.content
                })

                # Get next response from Claude
                response = self.anthropic.messages.create(
                    model="claude-3-5-sonnet-20241022",
                    max_tokens=1000,
                    messages=messages,
                )

                final_text.append(response.content[0].text)

        return "\n".join(final_text)

    async def chat_loop(self):
        """Run an interactive chat loop"""
        print("\nMCP Client Started!")
        print("Type your queries or 'quit' to exit.")
        
        while True:
            try:
                query = input("\nQuery: ").strip()
                
                if query.lower() == 'quit':
                    break
                    
                response = await self.process_query(query)
                print("\n" + response)
                    
            except Exception as e:
                print(f"\nError: {str(e)}")
    
    async def cleanup(self):
        """Clean up resources"""
        await self.exit_stack.aclose()

async def main():
    if len(sys.argv) < 2:
        print("Usage: python client.py <path_to_server_script>")
        sys.exit(1)
        
    client = MCPClient()
    try:
        await client.connect_to_server(sys.argv[1])
        await client.chat_loop()
    finally:
        await client.cleanup()

if __name__ == "__main__":
    import sys
    asyncio.run(main())